# Copyright (c) Alibaba, Inc. and its affiliates.
import os
from functools import partial
from typing import List, Union

from datasets import Dataset as HfDataset

from swift.plugin import extra_callbacks, get_loss_func, get_metric
from swift.trainers import TrainerFactory
from swift.utils import (append_to_jsonl, get_logger, get_model_parameter_info, is_master, plot_images, stat_array,
                         use_torchacc)
from ..argument import TrainArguments
from ..base import SwiftPipeline
from ..dataset import (EncodePreprocessor, GetLengthPreprocessor, IterablePackingDataset, LazyLLMDataset,
                       PackingDataset, load_dataset)
from ..infer import prepare_generation_config
from ..model import HfConfigFactory, get_model_arch
from ..utils import deep_getattr, dynamic_gradient_checkpointing
from .tuner import TunerMixin

logger = get_logger()


class SwiftSft(SwiftPipeline, TunerMixin):
    args_class = TrainArguments
    args: args_class

    def __init__(self, args: Union[List[str], TrainArguments, None] = None) -> None:
        logger.info(f"========1=====args:{args}")
        super().__init__(args)
        self.train_msg = {}
        self._prepare_model_tokenizer()
        self._prepare_template()
        self._prepare_callbacks()

    def _prepare_gradient_checkpointing(self):
        """配置梯度检查点(Gradient Checkpointing)以优化显存使用"""
        args = self.args

        # 禁用模型的缓存机制（与梯度检查点冲突）
        HfConfigFactory.set_model_config_attr(self.model, 'use_cache', False)

        # 如果启用梯度检查点配置
        if args.gradient_checkpointing:
            # 标记模型支持梯度检查点
            self.model.supports_gradient_checkpointing = True

            # 应用动态梯度检查点（按需选择检查点层）
            dynamic_gradient_checkpointing(self.model)

            # 启用输入梯度计算（某些特殊层需要）
            self.model.enable_input_require_grads()

        # 处理多模态模型的视觉塔部分
        model_meta = self.model.model_meta
        model_arch = get_model_arch(model_meta.model_arch)
        # 获取模型架构信息

        # 如果是多模态模型且有定义架构
        if model_meta.is_multimodal and model_arch:
            # 遍历所有视觉塔组件
            for vision_tower_name in model_arch.vision_tower:
                vision_tower = deep_getattr(self.model, vision_tower_name)

                # 如果视觉塔支持梯度需求配置
                if hasattr(vision_tower, 'enable_input_require_grads'):
                    try:
                        vision_tower.enable_input_require_grads()  # 启用输入梯度
                    except NotImplementedError:
                        pass  # 忽略不支持的组件

    def _prepare_generation_config(self):
        """准备生成配置，用于控制模型生成文本时的行为（如beam search、sampling等）"""
        args = self.args

        # 保存原始的generation_config以便后续恢复
        self.model.origin_generation_config = self.model.generation_config

        # 准备新的生成配置，合并默认配置和请求配置
        self.model.generation_config = prepare_generation_config(
            self.model.generation_config,
            # 模型默认的生成配置
            args.get_request_config(),
            # 从args获取的请求级别配置（可能包含用户自定义参数）
            self.tokenizer
            # 分词器（用于处理特殊token等）
        )

        # 记录最终使用的生成配置
        logger.info(f'model.generation_config: {self.model.generation_config}')

    def _prepare_model_tokenizer(self):
        # 获取训练参数
        args = self.args

        # 1. 如果启用了序列并行，初始化序列并行环境
        if args.sequence_parallel_size > 1:
            from swift.trainers.sequence_parallel import sequence_parallel
            sequence_parallel.init_sequence_parallel(args.sequence_parallel_size)

        # 2. 加载模型和分词器
        self.model, self.processor = args.get_model_processor()

        # 3. 如果模型使用了多 GPU 分配（hf_device_map），打印设备映射
        if hasattr(self.model, 'hf_device_map'):
            logger.info(f'model.hf_device_map: {self.model.hf_device_map}')

        # 4. 打印模型信息
        logger.info(f'model_info: {self.model.model_info}')

        # 5. 准备生成任务的配置
        self._prepare_generation_config()

        # 6. 准备梯度检查点（节省显存）
        self._prepare_gradient_checkpointing()

    def _prepare_template(self) -> None:
        """
        准备对话模板，用于格式化模型输入输出

        功能说明：
        1. 根据任务类型和处理器初始化模板
        2. 对因果语言模型(causal_lm)设置训练模式
        3. 如果模板需要使用模型本身，将当前模型赋值给模板
        """

        # 从args获取模板实例，并传入processor(处理器/分词器)
        template = self.args.get_template(self.processor)

        # 如果是因果语言模型任务，设置模板为训练模式
        # 训练模式通常会添加特殊token或调整输入格式
        if self.args.task_type == 'causal_lm':
            template.set_mode('train')

        # 如果模板需要访问模型(如获取embedding等)，将当前模型赋给模板
        if template.use_model:
            template.model = self.model

        # 将配置好的模板保存到实例变量中
        self.template = template

    def _get_dataset(self):
        """
        加载并准备训练集和验证集

        功能说明：
        1. 从配置参数加载主数据集，并自动划分训练/验证集
        2. 支持单独加载验证集（当val_dataset参数指定时）
        3. 提供数据集shuffle和比例划分功能
        """

        args = self.args

        # 获取数据集加载的额外参数（如预处理配置等）
        dataset_kwargs = args.get_dataset_kwargs()

        # 1. 加载主数据集（自动划分训练/验证集）
        train_dataset, val_dataset = load_dataset(
            # 主数据集名称/路径
            args.dataset,
            # 验证集划分比例（如0.1表示10%作为验证集）
            split_dataset_ratio=args.split_dataset_ratio,
            # 是否在划分前打乱数据
            shuffle=args.dataset_shuffle,
            # 其他数据集加载参数
            **dataset_kwargs
        )

        # 2. 处理独立验证集（当val_dataset参数非空时）
        if len(args.val_dataset) > 0:
            # 加载完整数据集作为验证集（split_dataset_ratio=1.0表示全部作为验证集）
            _, val_dataset = load_dataset(
                # 独立验证集路径
                args.val_dataset,
                # 使用全部数据
                split_dataset_ratio=1.0,
                # 是否打乱
                shuffle=args.val_dataset_shuffle,
                **dataset_kwargs
            )
            # 确保当使用独立验证集时，主数据集不进行划分
            assert args.split_dataset_ratio == 0., "当指定val_dataset时，split_dataset_ratio必须为0"

        # 记录数据集信息
        logger.info(f'train_dataset: {train_dataset}')
        logger.info(f'val_dataset: {val_dataset}')

        return train_dataset, val_dataset

    def _get_loss_func(self):
        args = self.args
        loss_type = args.loss_type
        if loss_type is None and args.loss_scale != 'default':
            loss_type = 'loss_scale'
        return get_loss_func(loss_type)

    def _get_data_collator(self):
        # 获取训练参数对象
        args = self.args

        # 获取当前使用的 Prompt 模板对象（用于生成输入格式）
        template = self.template

        # 如果训练类型是 longlora，则设置 padding_to 为最大长度（强制补全）
        # 否则 padding_to 为 None（即由 tokenizer 决定是否补全）
        padding_to = args.max_length if args.train_type == 'longlora' else None

        # 返回一个偏函数（partial），即带默认 padding_to 参数的 template.data_collator 函数
        return partial(template.data_collator, padding_to=padding_to)

    @staticmethod
    def _save_val_dataset(output_dir: str, val_dataset):
        if is_master() and isinstance(val_dataset, HfDataset):
            os.makedirs(output_dir, exist_ok=True)
            val_dataset_path = os.path.join(output_dir, 'val_dataset.jsonl')
            append_to_jsonl(val_dataset_path, val_dataset.to_list())
            logger.info(f'The split dataset from the training set will be saved at: {val_dataset_path}.')

    def run(self):
        args = self.args  # 获取实例中的参数配置

        # 获取训练集和验证集
        train_dataset, val_dataset = self._get_dataset()

        # 对数据集进行编码处理（例如 tokenizer 编码）
        train_dataset, val_dataset = self._encode_dataset(train_dataset, val_dataset)

        # 如果任务类型是序列分类
        if args.task_type == 'seq_cls':
            # 设置问题类型（如果未设置，则从模型配置中读取）
            args.problem_type = args.problem_type or getattr(self.model.config, 'problem_type', None)
            logger.info(f'args.problem_type: {args.problem_type}')  # 打印问题类型日志

        # 保存参数配置（可用于模型保存或恢复）
        args.save_args()

        # 获取数据整理器（Data Collator，用于按 batch 整理数据，特别是 padding 等）
        data_collator = self._get_data_collator()

        # 某些训练器（如 LoRA-GA）在准备阶段需要训练数据和 data_collator
        self.model = self.prepare_model(
            self.args, self.model, template=self.template, train_dataset=train_dataset
        )

        # 打印模型结构
        logger.info(f'model: {self.model}')

        # 获取模型参数信息（如参数总量、可训练参数量等）
        model_parameter_info = get_model_parameter_info(self.model)
        self.train_msg['model_parameter_info'] = model_parameter_info
        logger.info(f'model_parameter_info: {model_parameter_info}')

        # 获取训练器类（根据不同的训练任务或配置选择相应 Trainer 类）
        trainer_cls = TrainerFactory.get_trainer_cls(args)

        # 实例化 Trainer 对象
        trainer = trainer_cls(
            model=self.model,  # 模型对象
            args=self.args.training_args,  # 训练参数（如学习率、批大小等）
            data_collator=data_collator,  # 数据整理器
            train_dataset=train_dataset,  # 训练数据集
            eval_dataset=val_dataset,  # 验证数据集
            callbacks=self.callbacks,  # 回调函数（例如日志、早停等）
            template=self.template,  # 模板对象（如果任务需要模板，比如 prompt tuning）
            **self._get_trainer_kwargs(),  # 其它额外的 trainer 参数
        )

        # 调用训练函数开始训练过程
        return self.train(trainer)

    def _get_trainer_kwargs(self):
        args = self.args
        if args.metric is not None:
            compute_metrics, preprocess_logits_for_metrics = get_metric(args.metric)
        elif args.predict_with_generate:
            compute_metrics, preprocess_logits_for_metrics = get_metric('nlg')
        else:
            compute_metrics, preprocess_logits_for_metrics = get_metric('acc')
            compute_metrics = partial(
                compute_metrics, acc_strategy=args.acc_strategy, is_encoder_decoder=self.template.is_encoder_decoder)
        return {
            'compute_metrics': compute_metrics,
            'preprocess_logits_for_metrics': preprocess_logits_for_metrics,
            'compute_loss_func': self._get_loss_func()
        }

    def _save_trainer_state(self, trainer):
        training_args = trainer.args
        state = trainer.state
        if not hasattr(state, 'last_model_checkpoint'):
            logger.warning('No training was carried out, which may be due to the dataset being too small '
                           'or incorrect usage of resume_from_checkpoint.')
            return
        if self.args.create_checkpoint_symlink:
            last_checkpoint = os.path.join(self.args.output_dir, 'last')
            best_checkpoint = os.path.join(self.args.output_dir, 'best')
            os.symlink(state.last_model_checkpoint, last_checkpoint)
            os.symlink(state.best_model_checkpoint, best_checkpoint)
            state.last_model_checkpoint = last_checkpoint
            state.best_model_checkpoint = best_checkpoint
        logger.info(f'last_model_checkpoint: {state.last_model_checkpoint}')
        logger.info(f'best_model_checkpoint: {state.best_model_checkpoint}')

        # Visualization
        if is_master() and not use_torchacc():
            if 'tensorboard' in training_args.report_to:
                images_dir = os.path.join(training_args.output_dir, 'images')
                logger.info(f'images_dir: {images_dir}')
                plot_images(images_dir, training_args.logging_dir, ['train/loss'], 0.9)
            if training_args.push_to_hub:
                trainer.push_to_hub()

        self.train_msg.update({
            'last_model_checkpoint': state.last_model_checkpoint,
            'best_model_checkpoint': state.best_model_checkpoint,
            'best_metric': state.best_metric,
            'global_step': state.global_step,
            'log_history': state.log_history,
            'memory': trainer.max_memory,
        })
        if is_master():
            jsonl_path = os.path.join(training_args.output_dir, 'logging.jsonl')
            append_to_jsonl(jsonl_path, self.train_msg)
        return self.train_msg

    def train(self, trainer):
        logging_path = os.path.join(trainer.args.output_dir, 'logging.jsonl')
        logger.info(f'The logging file will be saved in: {logging_path}')
        trainer.train(trainer.args.resume_from_checkpoint)

        return self._save_trainer_state(trainer)

    def _prepare_callbacks(self):
        from .callback import DynamicLayerActivationCallback, TrainerAdapterCallback
        args = self.args
        callbacks = []
        if args.lisa_activated_layers > 0:
            assert args.train_type == 'full', 'LISA only supports full parameter training.'
            lisa_callback = DynamicLayerActivationCallback(
                n_layers=args.lisa_activated_layers,  # Number of layers to activate
                step_interval=args.lisa_step_interval,  # Step interval to update active layers
                model=self.model)
            lisa_callback.switch_active_layers()  # Make trainable parameters printing a correct value
            callbacks.append(lisa_callback)

        if args.is_adapter and args.train_type == 'adalora':
            callbacks.append(TrainerAdapterCallback(args))
        callbacks += extra_callbacks
        self.callbacks = callbacks

    def _stat_dataset(self, dataset: HfDataset):
        args = self.args
        if isinstance(dataset, HfDataset):
            dataset = GetLengthPreprocessor()(dataset, num_proc=args.dataset_num_proc)
            length = dataset['length']
        else:
            length = []
            for row in dataset:
                length.append(max([len(row[k]) for k in row.keys() if k.endswith('input_ids')]))
        _, stat_str = stat_array(length)
        logger.info(f'Dataset Token Length: {stat_str}')
        return stat_str

    def _encode_dataset(self, train_dataset, val_dataset):
        """
        对训练集和验证集进行编码处理

        参数:
            train_dataset: 训练数据集
            val_dataset: 验证数据集

        返回:
            处理后的训练集和验证集
        """
        # 获取模板对象
        template = self.template
        # 获取参数对象
        args = self.args
        # 获取输出目录，优先使用output_dir，否则使用save参数
        output_dir = getattr(args, 'output_dir', None) or getattr(args, 'save')
        # 保存验证数据集
        self._save_val_dataset(output_dir, val_dataset)
        # 检查是否是GRPO(一种RLHF类型)
        is_grpo = hasattr(args, 'rlhf_type') and args.rlhf_type == 'grpo'
        # 获取是否使用生成预测的标志
        predict_with_generate = getattr(args, 'predict_with_generate', False)

        # 如果不是GRPO类型
        logger.info(f"=========is_grpo:{is_grpo}")
        if not is_grpo:
            # 如果启用打包(packing)模式
            if args.packing:
                # 根据是否流式处理选择不同的打包数据集类
                packing_dataset_cls = IterablePackingDataset if args.streaming else PackingDataset
                # 对训练集进行打包处理
                train_dataset = packing_dataset_cls(
                    self.template, train_dataset, num_proc=args.dataset_num_proc, strict=args.strict)
                # 如果有验证集且不需要生成预测，则同样处理验证集
                print(f"................................................{train_dataset}")
                if val_dataset is not None:
                    val_dataset = packing_dataset_cls(
                        self.template, val_dataset, num_proc=args.dataset_num_proc, strict=args.strict)
            # 如果启用延迟(lazy)分词模式
            elif args.lazy_tokenize:
                # 使用LazyLLMDataset处理训练集
                train_dataset = LazyLLMDataset(
                    train_dataset, template.encode, strict=args.strict, random_state=args.data_seed)
                # 处理验证集(如果有且不需要生成预测)
                if val_dataset is not None and not predict_with_generate:
                    val_dataset = LazyLLMDataset(
                        val_dataset, template.encode, strict=args.strict, random_state=args.data_seed)
            # 如果不是打包也不是延迟分词模式
            else:
                # 使用编码预处理器
                preprocessor = EncodePreprocessor(template=template)
                # 预处理训练集
                train_dataset = preprocessor(train_dataset, num_proc=args.dataset_num_proc, strict=args.strict)
                # 预处理验证集(如果有且不需要生成预测)
                if val_dataset is not None and not predict_with_generate:
                    val_dataset = preprocessor(val_dataset, num_proc=args.dataset_num_proc, strict=args.strict)

            # 如果是主进程
            if is_master():
                # 获取第一个样本用于打印示例
                inputs = train_dataset[0] if hasattr(train_dataset, '__len__') else next(iter(train_dataset))
                # 打印输入示例
                template.print_inputs(inputs, tokenizer_kwargs=inputs.pop('tokenizer_kwargs', None) or {})
            # 如果是HfDataset或PackingDataset类型
            if isinstance(train_dataset, (HfDataset, PackingDataset)):
                # 统计训练集信息
                self.train_msg['train_dataset'] = self._stat_dataset(train_dataset)
                # 统计验证集信息(如果有且不需要生成预测)
                if val_dataset is not None and not predict_with_generate:
                    self.train_msg['val_dataset'] = self._stat_dataset(val_dataset)

        return train_dataset, val_dataset  # 返回处理后的数据集


def sft_main(args: Union[List[str], TrainArguments, None] = None):
    return SwiftSft(args).main()
